import base64
import datetime
from io import BytesIO
import numpy as np
import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import style

from sklearn.model_selection import cross_validate, cross_val_score, train_test_split
from sklearn import preprocessing, svm
from sklearn.linear_model import LinearRegression, BayesianRidge


def feature_engineering(context):
    data=context['data']
    data['HL_PCT'] = (data['high'] - data['low']) / data['close'] * 100.0
    data['PCT_change'] = (data['close'] - data['open']) / data['open'] * 100.0
    forecast_col = 'close'
    data.fillna(value=-float('inf'), inplace=True)
    forecast_out = int(np.math.ceil(0.01 * len(data)))
    data['label'] = data[forecast_col].shift(-forecast_out)
    context['data'] = data
    context['forecast_out'] = forecast_out
    return context


def fit_function(context):
    data = context['data']

    scale = context['scale']
    train_data = data[['close', 'HL_PCT', 'PCT_change', 'volume', 'label']]
    forecast_out = context['forecast_out']

    X = np.array(train_data.drop(['label'], 1))
    # X = preprocessing.scale(X)
    X_lately = X[-forecast_out:]

    X = X[:-forecast_out]
    train_data.dropna(inplace=True)
    y = np.array(train_data['label'])
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=scale)
    clf = BayesianRidge()
    clf.fit(X_train, y_train)
    accuracy = clf.score(X_test, y_test)
    context['accuracy'] = accuracy
    context['X_lately'] = X_lately
    context['clf'] = clf
    return context


def prediction_function(context):
    data = context['data']
    clf = context['clf']
    X_lately = context['X_lately']
    forecast_out = context['forecast_out']
    forecast_set = clf.predict(X_lately)
    data['Forecast'] = data['close']
    last_date = data.iloc[-1].date
    last_unix = last_date.timestamp()
    data.index = [x for x in data.date.values]
    # 预测forecast_out天后的
    one_day = 86400
    next_unix = last_unix + one_day
    # data['Forecast'] = np.nan
    for i in forecast_set:
        next_date = datetime.datetime.fromtimestamp(next_unix)
        next_unix += 86400
        data.loc[next_date] = [np.nan for _ in range(len(data.columns) - 1)] + [i]
    predict_data = data.tail(forecast_out + 1)

    context['predict_data'] = predict_data
    return context


def output_function(context):
    clf = context['clf']
    accuracy = context['accuracy']
    sources_data = context['data'][0:-context['forecast_out']]
    predict_data = context['predict_data']
    predict_data['date'] = predict_data.index
    forecast_data = predict_data[['date', 'Forecast']]
    forecast_data.index = [x for x in range(len(forecast_data))]
    del context
    return clf, accuracy, sources_data, forecast_data
